import json

from conf.neo4j_cnx import NeoCnxPool

from elasticsearch import Elasticsearch

from conf.variable import ES_HOSTS

es = Elasticsearch(hosts=ES_HOSTS)

neo_cnx = NeoCnxPool()


class IcdKg:

    @staticmethod
    def find_by_uid(uid, label):
        with neo_cnx.get_session() as session:
            query = (
                        'MATCH (n:%s {uid: $uid}) '
                        'RETURN n.properties AS properties, '
                        'n.child AS child'
                    ) % label
            results = session.run(query, uid=uid)
            for record in results:
                properties = json.loads(record['properties'])
                z = {**{'uid': uid, 'child': record['child']}, **properties}
                return z

    @staticmethod
    def find_child_uids(uid):
        with neo_cnx.get_session() as session:
            query = (
                'MATCH (a:ICD11 {uid: $uid}) '
                'RETURN a.child as child '
            )
            results = session.run(query, uid=uid)
            for record in results:
                return record['child']

    def find_child(self, uid, label):
        child = []
        uid_list = self.find_child_uids(uid)

        for i in uid_list:
            ch = self.find_by_uid(i, label)
            if ch is not None:
                child.append(ch)

        return child

    @staticmethod
    def find_parent(uid):
        queue = []

        def find_parent_recursion(uid):
            with neo_cnx.get_session() as session:
                query = (
                    'MATCH (a:ICD11 {uid: $uid}) '
                    'MATCH (a)-[:Parent]->(b) '
                    'RETURN b '
                    'LIMIT 1'
                )
                results = session.run(query, uid=uid)
                results = [record['b'] for record in results.data()]
                if len(results) > 0:
                    queue.append(results[0]['uid'])
                    find_parent_recursion(results[0]['uid'])

        find_parent_recursion(uid)
        return queue

    def fuzzy_search(self, index_term, label, class_kind='category', size=100):
        node_list = []

        results = es.search(index='cdss-icd11-index-term',
                            body={'size': size,
                                  'query': {
                                      'bool': {
                                          'must': [
                                              {'match': {'index_term': index_term}},
                                              {'term': {'class_kind': class_kind}}
                                          ]
                                      }
                                  }
                                  })

        for hit in results['hits']['hits']:
            uid = hit['_source']['uid']
            node = self.find_by_uid(uid, label)
            if node is not None:
                node_list.append(node)

        return node_list


if __name__ == '__main__':
    icd_kg_app = IcdKg()
    child = icd_kg_app.find_child(uid='1435254666', label='症候')

    child2 = icd_kg_app.find_child(uid='', label='疾病')

    results = icd_kg_app.fuzzy_search(index_term='咳嗽', label='症候')
    queue = icd_kg_app.find_parent(uid='1235618695')

    neo_cnx.close()
